"""
This is part of Human Robot Dynamics-Control, By Tourists.

Based on "OpenLoong Dynamics Control"
Transfer from mujoco to ISAAC SIM by IsaacLab

"""

#include <Eigen/Dense>
#include "data_bus.h"
import numpy as np

from ..common import DataBus
from ..common import LegState
from ..common import MotionState
from ..mathq import Bezier_1D

class FootPlacement():
    def __init__(self):
        self.kp_vx = 0.0
        self.kp_vy = 0.0
        self.kp_wz = 0.0
        self.legLength = 0.0
        self.stepHeight = 0.0
        self.phi = 0.0   ## phase varialbe for trajectory generation, must between 0 and 1
        self.tSwing = 0.0    ##swing time
        self.posStart_W = np.zeros(3) 
        self.posDes_W = np.zeros(3) 
        self.hipPos_W = np.zeros(3) 
        self.STPos_W = np.zeros(3)
        self.desV_W = np.zeros(3) 
        self.curV_W = np.zeros(3)
        self.desWz_W = 0.0
        self.base_pos = np.zeros(3)
        # double Trajectory(double phase, double des1, double des2);
        # void getSwingPos();
        # void dataBusRead(DataBus &robotState);
        # void dataBusWrite(DataBus &robotState);
        self.legState = LegState.LSt
        
        ### private
        self.pDesCur  = np.array([0.0, 0.0, 0.0])
        # self.__pDesCur  = np.array([0.0, 0.0, 0.0])
        self.yawCur = 0.0
        self.theta0 = 0.0
        self.omegaZ_W = 0.0
        self.hip_width = 0.0
        
    def Trajectory(self, phase, hei, len):
        Bswpid = Bezier_1D()
        para0 = 5.0
        para1 = 3.0
        for i in range(int(para0)):
            Bswpid.P.append(0.0)
        for i in range(int(para1)):
            Bswpid.P.append(1.0)

        output = 0.0
        if(self.phi<phase):
            output=hei*Bswpid.getOut(self.phi/phase)
        else:
            s=Bswpid.getOut((1.4-self.phi)/(1.4-phase))
            if(s>0):
                output=hei*s +len*(1.0-s)
            else:
                output=len
        return output
             
    def dataBusRead(self, robotState):
        self.posStart_W = robotState.swingStartPos_W
        self.desV_W = robotState.js_vel_des
        self.desWz_W = robotState.js_omega_des[2]
        # self.curV_W = robotState.dq.block<3,1>(0,0)
        self.curV_W = robotState.dq[0:3]
        self.phi = robotState.phi
        self.hipPos_W = robotState.posHip_W
        self.STPos_W = robotState.posST_W
        self.base_pos = robotState.base_pos
        self.tSwing  = robotState.tSwing
        self.theta0 = robotState.theta0
        self.yawCur = robotState.rpy[2]
        self.omegaZ_W = robotState.base_omega_W[2]
        self.hip_width = robotState.width_hips
        self.legState = robotState.legState
        
    def dataBusWrite(self, robotState):
        robotState.swingDesPosCur_W = np.array([self.pDesCur[0], self.pDesCur[1], self.pDesCur[2]])
        robotState.swingDesPosFinal_W = self.posDes_W
        robotState.swing_fe_rpy_des_W = np.array([0,0,robotState.base_rpy_des[2]]) ## WARNING! ThetaZ!
        robotState.swing_fe_pos_des_W = np.array([self.pDesCur[0], self.pDesCur[1], self.pDesCur[2]])

    def getSwingPos(self):
        b = np.zeros((4, 1))
        xNow = np.array([[1, self.phi, self.phi**2, self.phi**3]])

        KP = np.zeros((3, 3))
        KP[0, 0] = self.kp_vx
        KP[1, 1] = self.kp_vy
        KP[2, 2] = 0
        Rz = np.array([
            [np.cos(self.yawCur), -np.sin(self.yawCur), 0],
            [np.sin(self.yawCur), np.cos(self.yawCur), 0],
            [0, 0, 1]
        ])
        # KP=Rz*KP*Rz.transpose()
        KP = Rz @ KP @ Rz.T

        ## for linear velocity
        self.posDes_W = self.hipPos_W + KP @ (self.desV_W-self.curV_W) * (-1) + 0.5 * self.tSwing * self.curV_W + self.curV_W * (1 - self.phi) * self.tSwing

        ## for angular veloctity
        thetaF = 0.0
        thetaF = self.yawCur + self.theta0 + self.omegaZ_W * (1 - self.phi) * self.tSwing + 0.5 * self.omegaZ_W * self.tSwing + self.kp_wz * (self.omegaZ_W - self.desWz_W)
        self.posDes_W[0] += 0.5 * self.hip_width * (np.cos(thetaF) - np.cos(self.yawCur + self.theta0))
        self.posDes_W[1] += 0.5 * self.hip_width * (np.sin(thetaF) - np.sin(self.yawCur + self.theta0))

        xOff_L = -0.01 ## -0.01; // foot-end position offset in x direction in body frame
        yOff_L = 0.01 ## 0.01; // foot-end position offset in y direction in body frame, positive for moving the leg inside
        zOff_W = -0.035 ## foot-end position offset in z direction in world frame


        ## posDes_W(2)=STPos_W(2)-0.04;
        self.posDes_W[2] = self.base_pos[2] - self.legLength + zOff_W

        xOff_W = 0.0
        yOff_W = 0.0
        if self.legState == LegState.LSt:
            xOff_W = np.cos(self.yawCur) * xOff_L - np.sin(self.yawCur) * yOff_L
            yOff_W = np.sin(self.yawCur) * xOff_L + np.cos(self.yawCur) * yOff_L
        elif self.legState == LegState.RSt:
            xOff_W = np.cos(self.yawCur) * xOff_L - np.sin(self.yawCur) * (-yOff_L)
            yOff_W = np.sin(self.yawCur) * xOff_L + np.cos(self.yawCur) * (-yOff_L)


        self.posDes_W[0] += xOff_W
        self.posDes_W[1] += yOff_W
        
        #
        #    double yOff=0.005; // positive for moving the leg inside
        #    if (legState==DataBus::LSt)
        #        posDes_W(1)+=yOff;
        #    else if (legState==DataBus::RSt)
        #        posDes_W(1)-=yOff;

        ## cycloid trajectories
        if self.phi < 1.0:
            self.pDesCur[0] = self.posStart_W[0] + (self.posDes_W[0] - self.posStart_W[0]) / (2 * np.pi) * (2 * np.pi * self.phi - np.sin(2 * np.pi * self.phi))
            self.pDesCur[1] = self.posStart_W[1] + (self.posDes_W[1] - self.posStart_W[1]) / (2 * np.pi) * (2 * np.pi * self.phi - np.sin(2 * np.pi * self.phi))
            
        # self.pDesCur[2]=posStart_W(2)+self.stepHeight*0.5*(1-np.cos(2*3.1415*self.phi))+(self.posDes_W(2)-self.posStart_W(2))/(2*3.1415)*(2*3.1415*phi-sin(2*3.1415*phi));
        self.pDesCur[2] = self.posStart_W[2] + self.Trajectory(0.2, self.stepHeight, self.posDes_W[2] - self.posStart_W[2])



